import pandas as pd
import numpy as np
from math import log2

# 定义数据集
data = {
    'RID': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
    'age': ['youth', 'youth', 'middle_aged', 'senior', 'senior', 'senior', 'middle_aged', 'youth', 'youth', 'senior',
            'youth', 'middle_aged', 'middle_aged', 'senior'],
    'income': ['high', 'high', 'high', 'medium', 'low', 'low', 'low', 'medium', 'low', 'medium', 'medium', 'medium',
               'high', 'medium'],
    'student': ['no', 'no', 'no', 'no', 'yes', 'yes', 'yes', 'no', 'yes', 'yes', 'yes', 'no', 'yes', 'no'],
    'credit_rating': ['fair', 'excellent', 'fair', 'fair', 'fair', 'excellent', 'excellent', 'fair', 'fair', 'fair',
                      'excellent', 'excellent', 'fair', 'excellent'],
    'buys_computer': ['no', 'no', 'yes', 'yes', 'yes', 'no', 'yes', 'no', 'yes', 'yes', 'yes', 'yes', 'yes', 'no']
}
df = pd.DataFrame(data)


def calculate_entropy(data_column):
    """
    计算给定数据列（Pandas Series）的熵。
    :param data_column: 包含类别标签的Pandas Series。
    :return: 熵值。
    """
    value_counts = data_column.value_counts()
    total_instances = len(data_column)
    entropy = 0
    if total_instances == 0:  # 处理空子集的情况
        return 0
    for count in value_counts:
        probability = count / total_instances
        if probability > 0:  # 避免log2(0)
            entropy -= probability * log2(probability)
    return entropy


def calculate_information_gain(data, attribute_name, target_name):
    """
    计算指定属性相对于目标变量的信息增益。
    :param data: 包含数据的Pandas DataFrame。
    :param attribute_name: 要计算信息增益的属性列名。
    :param target_name: 目标变量的列名。
    :return: 信息增益值。
    """
    total_entropy = calculate_entropy(data[target_name])

    weighted_entropy_of_subsets = 0
    unique_values = data[attribute_name].unique()
    total_instances = len(data)

    if total_instances == 0:  # 处理空数据集的情况
        return 0

    for value in unique_values:
        subset = data[data[attribute_name] == value]
        subset_entropy = calculate_entropy(subset[target_name])
        proportion_of_subset = len(subset) / total_instances
        weighted_entropy_of_subsets += proportion_of_subset * subset_entropy

    information_gain = total_entropy - weighted_entropy_of_subsets
    return information_gain


def id3(data, original_features, target_attribute_name):
    """
    ID3算法的简化概念性实现。
    :param data: 当前节点的数据子集 (Pandas DataFrame)。
    :param original_features: 原始特征列表。
    :param target_attribute_name: 目标属性的名称。
    :return: 一个表示决策树（或子树）的字典，或者一个类别标签（叶节点）。
    """
    # 1. 如果当前子集所有实例属于同一类别，则返回该类别标签 (纯节点)
    if len(data[target_attribute_name].unique()) == 1:
        return data[target_attribute_name].iloc[0]  # 返回唯一的类别

    # 2. 如果没有特征可以用于分裂，或者数据为空，则返回数据集中最常见的类别标签
    if len(original_features) == 0 or len(data) == 0:
        return data[target_attribute_name].mode().iloc[0] if not data.empty else "Unknown"  # 处理空数据的情况

    # 3. 选择信息增益最高的特征作为当前节点的分裂特征
    gains = {feature: calculate_information_gain(data, feature, target_attribute_name)
             for feature in original_features}

    # 如果所有增益都为0或负（可能由于数据特性或浮点精度），则返回多数类
    if not gains or max(gains.values()) <= 0:
        return data[target_attribute_name].mode().iloc[0]

    best_feature = max(gains, key=gains.get)

    # 4. 创建树结构
    tree = {best_feature: {}}

    # 5. 移除已选择的特征，为递归调用准备新的特征集
    remaining_features = [f for f in original_features if f != best_feature]

    # 6. 为最佳特征的每个值创建一个分支，并递归调用id3
    for value in data[best_feature].unique():
        subset_data = data[data[best_feature] == value].copy()  # 使用 .copy() 避免 SettingWithCopyWarning

        # 如果子集为空，则分配父节点中最常见的类别
        if len(subset_data) == 0:
            tree[best_feature][value] = data[target_attribute_name].mode().iloc[0]
        else:
            tree[best_feature][value] = id3(subset_data, remaining_features, target_attribute_name)

    return tree


# --- 脚本执行部分 ---
print("原始数据集 df:")
print(df)
print("-" * 30)

# 示例：计算整个数据集目标变量的熵
initial_entropy = calculate_entropy(df['buys_computer'])
print(f"初始熵 (整个数据集): {initial_entropy:.4f}")
print("-" * 30)

# 示例：计算各属性的信息增益
features_list = ['age', 'income', 'student', 'credit_rating']
print("各属性对于初始数据集的信息增益:")
for feature in features_list:
    gain = calculate_information_gain(df, feature, 'buys_computer')
    print(f"  属性 '{feature}': {gain:.4f}")
print("-" * 30)

# 构建决策树
decision_tree_result = id3(df.copy(), features_list, 'buys_computer')  # 使用 df.copy() 确保原始df不被修改
print("构建的决策树 (字典表示):")
import json  # 导入json模块以美化打印字典

print(json.dumps(decision_tree_result, indent=4, ensure_ascii=False))
print("-" * 30)


# 预测新顾客 [senior, low, yes, fair]
def predict(tree, sample):
    """
    使用构建好的决策树进行预测。
    :param tree: 决策树（嵌套字典）。
    :param sample: 新样本（字典，键为属性名，值为属性值）。
    :return: 预测的类别标签。
    """
    # 如果当前tree是叶节点（直接是类别字符串）
    if not isinstance(tree, dict):
        return tree

    # 获取当前节点的测试属性
    attribute = next(iter(tree))

    # 获取样本在该属性上的值
    attribute_value = sample.get(attribute)

    # 如果样本中没有该属性值，或者该属性值在树的分支中不存在
    if attribute_value is None or attribute_value not in tree[attribute]:
        # 简单处理：可以返回一个默认值或尝试多数类等，这里返回"Unknown"
        # 更复杂的处理可能涉及查找父节点的多数类
        print(f"警告: 属性 '{attribute}' 的值 '{attribute_value}' 在决策树分支中未找到。")
        # 尝试找到一个默认分支或最常见分支，这里简化处理
        # 实际应用中，可能需要更鲁棒的未知值处理策略
        # 例如，可以尝试在当前节点下选择多数类，或者返回一个特定的“无法预测”标签
        # 为了简单起见，如果分支不存在，我们尝试返回当前节点下最常见的结果（如果可以确定）
        # 或者，如果无法确定，则返回一个通用标签
        # 这是一个简化的处理，实际应用中可能需要更复杂的逻辑
        # 尝试从树的当前层级中找到一个默认的叶节点值
        for key in tree[attribute]:
            if not isinstance(tree[attribute][key], dict):  # 找到一个叶节点
                return tree[attribute][key]
        return "无法预测 (路径不存在)"

    # 递归地沿着树向下
    return predict(tree[attribute][attribute_value], sample)


new_customer = {'age': 'senior', 'income': 'low', 'student': 'yes', 'credit_rating': 'fair'}
prediction = predict(decision_tree_result, new_customer)
print(f"对新顾客 {new_customer} 的预测结果: buys_computer = {prediction}")

